# -*- coding:utf-8 -*-
from common.log import log
from common.database_tools import MongoTools


def assist_handle_mongo_data(*args, **kwargs):
    """
    批量获取mongo数据
    :param args: args【0】迭代函数，args【1】：配置文件
    :param kwargs: args:额外参数，find 过滤条件,limit 单次获取长度,is_count 是否查询总条数进行分页
    :return:
    """

    conf = args[1]
    table = conf["table"]
    connect = conf["connect"]
    db = conf["db"]
    arg = kwargs.get("args", list())
    find = kwargs.get("find", None)
    limit = kwargs.get("limit", 600000)
    is_count = kwargs.get("is_count", True)
    mtl = MongoTools(connect, db)
    if is_count:
        count = mtl.get_count(find, table=table)
        log.info(f"获取总条数：{count},当前翻页条数:{limit}")
        for ct in range((count + limit - 1) // limit):
            skip = limit * ct
            log.info(f"总条数：{count}，目前查询范围：{skip}~{skip + limit}")
            datas = mtl.get_data(find, limit=limit, skip=skip, table=table)
            args[0](datas, *arg)
    else:
        datas = mtl.get_data(find, table=table)
        args[0](datas, *arg)
    del mtl


# get_aggregate_data
def aggregate_single_data(st, table, connect):
    """
    聚合单个字段，value
    :param st:
    :param table:
    :param connect:
    :return: list
    """
    mgl = MongoTools(connect)
    datas = mgl.get_aggregate_data([{"$match": {"company_host": {"$regex": "成都"}}}, {"$group": {"_id": f"${st}"}}],
                                   table)
    del mgl
    return [dt["_id"] for dt in datas]
